import numpy as np
import random
import os
import pandas as pd
node_cpu_capacity = 100
node_mem_capacity = 100
node_disk_capacity = 100


def get_datas_dir():
    return os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))+"/data_process/dealed_data"

def get_ms_number_time_cover(plan_file_path,ms_kind_list):
    # data_dir=get_datas_dir()
    data=pd.read_csv(plan_file_path, index_col=0)
    ms_number_list=[]
    for ms_kind in ms_kind_list:
        try:
            ms_number=round(data.loc[ms_kind].mean())
        except:
            ms_number=0
        ms_number_list.append(ms_number)
    # print(f"部署总数：{sum(ms_number_list)}")
    return ms_number_list

def get_ms_number_each_one(ms_ratio, all_number):
    a = np.asarray(ms_ratio, dtype=float)
    n = a.size
    k = int(all_number)

    # 规范化 k 的范围
    if k <= 0:
        return [0] * n
    if k >= n:
        return [1] * n

    # 取出前 k 大的索引（集合正确，顺序不保证）
    idx = np.argpartition(a, -k)[-k:]

    mask = np.zeros(n, dtype=int)
    mask[idx] = 1
    
    return mask

def get_ms_number_random(ms_ratio, all_number):
    n=len(ms_ratio)
    mask = [0] * n
    
    # 随机选择m个索引位置，设为1
    selected_indices = random.sample(range(n), all_number)
    for idx in selected_indices:
        mask[idx] = 1
    
    return mask

def get_ms_number_frequency(ms_ratio, all_number):
    # ms_number_ratio = ms_ratio / ms_ratio.sum()  # 归一化
    
    # ms_number = np.round(ms_number_ratio * all_number).astype(int)
    
    ms_number_ratio = np.asarray(ms_ratio, dtype=float)
    
    s = ms_number_ratio.sum()
    if s == 0:
        # 比例全为 0：无法区分分配，统一给 0（或根据你的业务改成把全部给第一个等）
        print("ms_ratio wrong (zero)!")
        exit(-1)

    # 归一化比例，确保总和为 1
    p = ms_number_ratio / s

    # 实数目标值
    raw = p * int(all_number)

    # 下整，保证不超过总量
    ms_number = np.floor(raw).astype(int)

    # 还需补的名额
    diff = int(all_number) - ms_number.sum()
    if diff > 0:
        # 取小数部分最大的 diff 个加 1（使用 argpartition 线性时间选择）
        remainder = raw - ms_number
        top_idx = np.argpartition(-remainder, diff - 1)[:diff]
        ms_number[top_idx] += 1

    # 此时必然 sum(ms_number) == all_number
    
    return ms_number


def Heuristic_algorithm(plan_file_path, ms_kind_list, weight_matrix, cpu_cost, mem_cost, disk_cost, ms_ratio, node_num,  resource_cost_ratio, heuristic_deploy):
    ms_number_list = get_ms_number_time_cover(plan_file_path,ms_kind_list)
    number=np.array(ms_number_list).sum()
    if heuristic_deploy=="time_cover":
        pass
    elif heuristic_deploy=="each_one":
        ms_number_list = get_ms_number_each_one(ms_ratio,number)
    elif heuristic_deploy=="frequency":
        ms_number_list = get_ms_number_frequency(ms_ratio,number)
    elif heuristic_deploy=="random":
        ms_number_list = get_ms_number_random(ms_ratio,number)
    else:
        print(f"heuristic_deploy wrong ({heuristic_deploy})!")
        exit(-1)
    
    # ms_types = len(ms_number_list)

    # 1. 初始化每个节点的剩余资源
    node_resources = [
        {'cpu': node_cpu_capacity, 'mem': node_mem_capacity, 'disk': node_disk_capacity}
        for _ in range(node_num)
    ]

    # 2. 记录部署结果
    ms_type_result = []  # 微服务类型ID
    node_id_result = []  # 节点ID

    # 3. 先将数量最多的微服务分散部署
    ms_left = ms_number_list.copy()
    ms_order = np.argsort(-np.array(ms_left))  # 数量从多到少的微服务类型index
    # 只部署数量最多的微服务
    ms_type = int(ms_order[0])
    count = ms_left[ms_type]
    node_idx = 0
    while count > 0:
        if (node_resources[node_idx]['cpu'] >= cpu_cost[ms_type] and
            node_resources[node_idx]['mem'] >= mem_cost[ms_type] and
            node_resources[node_idx]['disk'] >= disk_cost[ms_type]):
            node_resources[node_idx]['cpu'] -= cpu_cost[ms_type]
            node_resources[node_idx]['mem'] -= mem_cost[ms_type]
            node_resources[node_idx]['disk'] -= disk_cost[ms_type]
            ms_type_result.append(ms_type)
            node_id_result.append(node_idx)
            count -= 1
        node_idx = (node_idx + 1) % node_num
        if node_idx == 0 and count > 0 and all(
            node_resources[n]['cpu'] < cpu_cost[ms_type] or
            node_resources[n]['mem'] < mem_cost[ms_type] or
            node_resources[n]['disk'] < disk_cost[ms_type]
            for n in range(node_num)
        ):
            break
    assert count == 0
    ms_left[ms_type] = count

    # 4. 基于调用矩阵分配剩余微服务
    node_ms_types = [[] for _ in range(node_num)]
    for ms_type0, node_id in zip(ms_type_result, node_id_result):
        node_ms_types[node_id].append(ms_type0)

    # 只要还有剩余微服务就继续分配
    while sum(ms_left) > 0:
        ms_type = int(np.argmax(ms_left))
        if ms_left[ms_type] == 0:
            break
        N = ms_left[ms_type]
        node_call_score = np.zeros(node_num)
        for node_id in range(node_num):
            for deployed_ms in node_ms_types[node_id]:
                node_call_score[node_id] += weight_matrix[deployed_ms][ms_type]
        if np.all(node_call_score == 0):
            # 均匀分散部署，优先选择资源剩余量大的节点
            node_resource_score = np.array([
                node_resources[n]['cpu'] + node_resources[n]['mem'] + node_resources[n]['disk']
                for n in range(node_num)
            ])
            candidate_nodes = np.argsort(-node_resource_score)
            alloc = np.zeros(node_num, dtype=int)
            for i in range(N):
                alloc[candidate_nodes[i % node_num]] += 1
        else:
            # 按调用需求比例分配
            total_score = node_call_score.sum()
            if total_score == 0:
                alloc = np.zeros(node_num, dtype=int)
            else:
                alloc = np.round(N * node_call_score / total_score).astype(int)
            # 修正分配数量总和
            diff = N - alloc.sum()
            if diff > 0:
                # 按照最大余数法补齐
                remainders = N * node_call_score / total_score - alloc
                for idx in np.argsort(-remainders)[:diff]:
                    alloc[idx] += 1
            elif diff < 0:
                # 超分配则按最小余数法减少
                remainders = alloc - N * node_call_score / total_score
                for idx in np.argsort(-remainders)[:abs(diff)]:
                    alloc[idx] -= 1
        # 依次尝试部署
        deployed_count = 0
        for node_id in np.argsort(-alloc):
            for _ in range(alloc[node_id]):
                if ms_left[ms_type] == 0:
                    break
                if (node_resources[node_id]['cpu'] >= cpu_cost[ms_type] and
                    node_resources[node_id]['mem'] >= mem_cost[ms_type] and
                    node_resources[node_id]['disk'] >= disk_cost[ms_type]):
                    node_resources[node_id]['cpu'] -= cpu_cost[ms_type]
                    node_resources[node_id]['mem'] -= mem_cost[ms_type]
                    node_resources[node_id]['disk'] -= disk_cost[ms_type]
                    ms_type_result.append(ms_type)
                    node_id_result.append(int(node_id))
                    node_ms_types[node_id].append(ms_type)
                    ms_left[ms_type] -= 1
                    deployed_count += 1
        # 若有未分配完的，按资源剩余量大的节点补齐
        if ms_left[ms_type] > 0:
            node_resource_score = np.array([
                node_resources[n]['cpu'] + node_resources[n]['mem'] + node_resources[n]['disk']
                for n in range(node_num)
            ])
            candidate_nodes = np.argsort(-node_resource_score)
            for node_id in candidate_nodes:
                while ms_left[ms_type] > 0 and \
                    node_resources[node_id]['cpu'] >= cpu_cost[ms_type] and \
                    node_resources[node_id]['mem'] >= mem_cost[ms_type] and \
                    node_resources[node_id]['disk'] >= disk_cost[ms_type]:
                    node_resources[node_id]['cpu'] -= cpu_cost[ms_type]
                    node_resources[node_id]['mem'] -= mem_cost[ms_type]
                    node_resources[node_id]['disk'] -= disk_cost[ms_type]
                    ms_type_result.append(ms_type)
                    node_id_result.append(node_id)
                    node_ms_types[node_id].append(ms_type)
                    ms_left[ms_type] -= 1
                if ms_left[ms_type] == 0:
                    break
    return ms_type_result, node_id_result


if __name__ == "__main__":
    # 微服务权重矩阵
    weight_matrix = np.array([[0, 1,0],[0,0,0], [6, 7,0]])
    cpu_cost = np.array([2, 5,1])
    mem_cost = np.array([3, 4,2])
    disk_cost = np.array([5, 6,2])
    ms_ratio = np.array([0.3, 0.2,0.1])

    node_num = 4
    
    resource_cost_ratio=0.5
    

    ms_types = len(cpu_cost)
    
    # print("每类微服务数量:", ms_number)
    ms_type_result, node_id_result = Heuristic_algorithm(weight_matrix, cpu_cost, mem_cost, disk_cost, ms_ratio, node_num, resource_cost_ratio)
    print("微服务类型ID:", ms_type_result)
    print("节点ID:", node_id_result)


